# -*-coding:utf-8-*-
"""
二值化，必须保存为png图片，jpg图片并不是黑白图片
"""

import sys, os
import time
import pytesseract
from PIL import Image

# reload(sys)
# sys.setdefaultencoding('utf-8')
# time1 = time.time()


def binarizing(img, threshold):
    pixdata = img.load()
    w, h = img.size
    for y in range(h):
        for x in range(w):
            if pixdata[x, y] < threshold:
                pixdata[x, y] = 0
            else:
                pixdata[x, y] = 255
    return img


###########去除干扰线算法
def depoint(img):  # input: gray image
    pixdata = img.load()
    w, h = img.size
    for y in range(1, h - 1):
        for x in range(1, w - 1):
            count = 0
            if pixdata[x, y - 1] < 111:
                count = count + 1
            if pixdata[x, y + 1] < 111:
                count = count + 1
            if pixdata[x - 1, y] < 111:
                count = count + 1
            if pixdata[x + 1, y] < 111:
                count = count + 1
            if count > 2:
                pixdata[x, y] = 0

    return img


def test_ocr(file_name):
    image = Image.open(file_name)
    # img_ = image.convert('L')
    # img_ = depoint(img_)
    # img_ = binarizing(img_,150)

    # img_.show()
    code = pytesseract.image_to_string(image)
    # print str(code)


def test_image_repair(filename):
    image = Image.open(filename)
    img_ = image.convert('L')
    # img_ = depoint(img_)
    img_ = binarizing(img_, 100)
    img_.show()
    code = pytesseract.image_to_string(img_, config="-psm 7", nice=1)
    ## print str(code)
    ## print pix


def listDir(rootDir):
    file_list = []
    for filename in os.listdir(rootDir):
        pathname = os.path.join(rootDir, filename)
        if (os.path.isfile(pathname)):
            file_list.append(pathname)
        else:
            file_list.append(listDir(pathname))
    return file_list


def test_ocr1():
    file_list = []
    file_list = listDir("ocr_image")
    # print file_list
    for file_name in file_list:
        test_ocr(file_name)
        break


def test_func():
    a = 110

    def test_func2():
        b = a
        # print b

    t = test_func2
    test_func2()
    # print a


def do_binarizing():
    pth = r'N:\_TMP\captcha\battle_net_result_binarizing'
    os.chdir(pth)

    file_list = listDir(pth)
    for file_name in file_list:
        image = Image.open(file_name)
        img_ = image.convert('L')
        # img_ = depoint(img_)
        img_ = binarizing(img_, 150)
        # file_name = file_name.replace('TEST_DATA','TEST_DATA_NEW')
        # file_name = file_name.replace(".do",".png")
        img_.save(file_name)
        # print file_name


def binarizing2(img, threshold):
    table = []
    for i in range(256):
        if i < threshold:
            table.append(0)
        else:
            table.append(1)

    return img.point(table, '1')


def do_binarizing_file(pth_src, pth_dst):
    image = Image.open(pth_src)
    img_ = image.convert('L')
    img2 = binarizing2(img_, 210)
    img2.save(pth_dst)


def do_binarizing2(pth_src, pth_dst):
    os.chdir(pth_src)

    file_list = os.listdir(pth_src)
    for file_name in file_list:
        print(file_name)
        image = Image.open(file_name)
        img_ = image.convert('L')
        img2 = binarizing2(img_, 210)
        img2.save(pth_dst + '\\' + file_name)
        # return


if __name__ == "__main__":
    # test_func()
    # test_image_repair('ocr_image\\3.png')
    # test_ocr1()
    pth_src = r'N:\_TMP\captcha\battle_net_result'
    pth_dst = r'N:\_TMP\captcha\battle_net_result_binarizing_test'

    pth_src = r'N:\_TMP\captcha_prj_bn_lsmt\predict_testing_binarizing'
    do_binarizing2(pth_src, pth_src)
